import json
import requests
import random
import logging
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
import queue

input_file = '0721/pre-train-texts.jsonl'
output_file = '0721/finetune-data.jsonl'
log_file = '0721/generate-finetune-data.log'

# 日志配置
logging.basicConfig(
    filename=log_file,
    filemode='a',
    format='%(asctime)s %(levelname)s: %(message)s',
    level=logging.INFO
)

model_services = [
    {
        "api_url": "http://18.0.31.1:9981/v1/chat/completions",
        "model": "ds671"
    },
    {
        "api_url": "http://18.0.31.2:9981/v1/chat/completions",
        "model": "/home/llm/qwq32"
    }
]

# instruction模板
instructions = [
    # 通用类
    "根据规章制度回答以下问题",
    "请依据提供的制度条款作答",
    "请根据公司规章内容回答以下问题",
    "根据相关管理规定，回答以下问题",
    # 任务导向类
    "请从制度文本中提取相关信息回答问题",
    "请根据提供的条款内容进行回答",
    "请依据制度内容给出准确回答",
    # 指令风格类
    "请回答以下与制度相关的问题",
    "请根据给定的规章制度内容作出回答",
    "请从制度文本中找出答案",
    # 专业术语类
    "请根据制度条款解释以下内容",
    "请结合制度内容说明以下问题",
    "请依据管理规定对以下问题进行解释"
]

# 构造大模型prompt
def build_prompt(text):
    return f"""你是一个专业助手，负责根据提供的规章制度文本，生成用于微调训练的问答对。每条原始文本需生成至少3条不同的问答对，格式如下：

{{
  "instruction": "根据规章制度回答以下问题",
  "input": "问题内容",
  "output": "答案内容"
}}

请根据以下原始文本内容，生成3个以上不同角度的问答对：

[原始文本]
{text}

要求如下：
1. 问题（input）应基于文本内容，合理且具体，不能太泛或模糊。
2. 答案（output）必须严格基于原文内容，不能编造或扩展。
3. 问题可以包括定义解释、规则适用、条款内容、流程说明等。
4. 输出为标准JSON格式的问答对列表，每条包含instruction、input、output三个字段。
5. 不使用Markdown格式，不添加额外说明，只输出结果。
其中 instruction可以在以下内容中随机选择：
{instructions}
"""

def call_openai_api(api_url, model, prompt):
    headers = {"Content-Type": "application/json"}
    data = {
        "model": model,
        "messages": [
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.7
    }
    try:
        resp = requests.post(api_url, headers=headers, json=data, timeout=60)
        resp.raise_for_status()
        result = resp.json()
        content = result['choices'][0]['message']['content'].strip()
        logging.info(f"API返回内容: {content}")
        return content
    except Exception as e:
        logging.error(f"API调用失败: {e}")
        return None

def parse_qa_list(response_text):
    """
    解析大模型返回的JSON列表，失败时记录日志
    """
    try:
        # 兼容模型返回前后有多余字符的情况
        start = response_text.find('[')
        end = response_text.rfind(']')
        if start != -1 and end != -1:
            json_str = response_text[start:end+1]
            qa_list = json.loads(json_str)
            return [
                {
                    "instruction": qa.get("instruction", "").strip(),
                    "input": qa.get("input", "").strip(),
                    "output": qa.get("output", "").strip()
                }
                for qa in qa_list
                if "instruction" in qa and "input" in qa and "output" in qa
            ]
        else:
            logging.error(f"未找到JSON列表，原始返回: {response_text}")
            return []
    except Exception as e:
        logging.error(f"解析模型返回失败: {e}，原始返回: {response_text}")
        return []

def process_one(idx, text):
    results = []
    prompt = build_prompt(text)
    for service in model_services:
        response = call_openai_api(service["api_url"], service["model"], prompt)
        if response:
            qa_list = parse_qa_list(response)
            if not qa_list:
                logging.warning(f"第{idx+1}行，模型{service['model']}未解析出Q&A，原始返回: {response}")
            for qa in qa_list:
                qa["instruction"] = random.choice(instructions)
                results.append(json.dumps(qa, ensure_ascii=False))
        else:
            logging.error(f"第{idx+1}行，模型{service['model']}无返回")
    return results

def main():
    # 线程安全队列
    result_queue = queue.Queue()
    lock = threading.Lock()
    lines = []
    with open(input_file, 'r', encoding='utf-8') as fin:
        for line in fin:
            line = line.strip()
            if line:
                lines.append(line)

    def worker(idx, line):
        try:
            item = json.loads(line)
            text = item.get("text", "")
        except Exception as e:
            logging.error(f"第{idx+1}行解析失败: {e}，内容: {line}")
            return
        results = process_one(idx, text)
        for r in results:
            result_queue.put(r)

    max_workers = min(16, len(lines))  # 可根据实际情况调整线程数
    with ThreadPoolExecutor(max_workers=max_workers) as executor, \
            open(output_file, 'w', encoding='utf-8') as fout:
        futures = []
        for idx, line in enumerate(lines):
            futures.append(executor.submit(worker, idx, line))
        # 实时写入，避免内存占用过大
        finished = 0
        total = len(futures)
        with tqdm(total=total, desc="处理中") as pbar:
            while finished < total:
                try:
                    r = result_queue.get(timeout=1)
                    with lock:
                        fout.write(r + "\n")
                except queue.Empty:
                    pass
                # 检查已完成的future
                finished = sum(1 for f in futures if f.done())
                pbar.n = finished
                pbar.refresh()
        # 写入剩余结果
        while not result_queue.empty():
            r = result_queue.get()
            with lock:
                fout.write(r + "\n")

if __name__ == "__main__":
    main()
